from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,#加载因果模型
    DataCollatorForLanguageModeling,#数据集成器
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset

#设置预训练模型的路径
model_name = "/root/lanyun-tmp/models/DeepSeek-R1-Distill-Qwen-1.5B"
#设置微调后的模型保存路径
finetuned_model_path = "/root/lanyun-tmp/models/new_model"
#加载分词器
tokenizer = AutoTokenizer.from_pretrained(model_name)
#加载预训练模型
model = AutoModelForCausalLM.from_pretrained(model_name,local_files_only=True)#加载本地模型
#加载数据集
train_prompt_style = '''
以下是描述任务的说明，并附有提供进一步背景的输入。
请写一个恰当的回答，完成该任务。
在回答之前，请仔细思考问题，并创建一个逐步的思维链，以确保逻辑清晰且准确。
### 任务描述:
你是一位在临床推理，诊断和治疗计划方面具有高级知识的医学专家。
请回答以下医学问题。

### 问题:
{}

### 回答:
<think>
{}
<think>
{}
'''
EOS_TOKEN = tokenizer.eos_token
#数据集预处理函数
def process_data(examples):
    inputs = examples["Question"]
    cots = examples["Complex_CoT"]
    outputs = examples["Response"]
    texts = []
    for inputs, cot, outputs in zip(inputs, cots, outputs):
        text = train_prompt_style.format(inputs, cot, outputs) + EOS_TOKEN
        texts.append(text)
    return tokenizer(texts,max_length=2048,truncation=True)

dataset = load_dataset("/root/lanyun-tmp/datasets/medical-o1-reasoning-SFT","zh",split="train[0:2000]")
dataset = dataset.map(process_data, batched=True, remove_columns=dataset.column_names)

# print(model)
#设置模型微调的参数
from peft import LoraConfig, get_peft_model
from peft.utils import TaskType
#LoRA配置
lora_config = LoraConfig(
    r=16,#秩
    lora_alpha=32,#缩放参数 2*r
    target_modules=['self_attn.q_proj','self_attn.k_proj','self_attn.v_proj','self_attn.o_proj'],
    lora_dropout = 0.05,
    bias='none',
    task_type = TaskType.CAUSAL_LM,
)
model = get_peft_model(model, lora_config)

#实际上直接计算所有的可计算参数
def print_trainable_parameters(model):
    trainable_params = 0
    all_params = 0
    for _, param in model.named_parameters():
        num_param = param.numel()
        all_params += num_param
        if param.requires_grad:
            trainable_params += num_param
    print(f"Total params: {all_params}, Trainable params: {trainable_params}")

print_trainable_parameters(model)

model.enable_input_require_grads() #访问梯度检查点的时候要执行这个方法
#配置训练参数
import os
os.environ["SWANLAB_PROJECT"] = "test"
os.environ["SWANLAB_WORKSPACE"] = "XYLCURRY30"
# train_args = TrainingArguments(
#     output_dir=finetuned_model_path, #输出文件夹
#     num_train_epochs=100, #训练轮数
#     overwrite_output_dir=True, #覆盖输出目录的现有内容
#     per_device_train_batch_size=1, #训练批次大小
#     save_steps=100, #每隔多少步保存一次模型
#     save_total_limit=20, #保存的模型文件数量限制
#     prediction_loss_only=True, #在评估时只计算损失
#     report_to="swanlab", #启用swanlab报告
#     logging_dir="./logs",
#     save_strategy="epoch",
#     gradient_accumulation_steps=16, #梯度累积步数, 即每多少步更新一次梯度
#     gradient_checkpointing=True, #使用梯度检查点以节省内存
#     remove_unused_columns=False, #是否移除未使用的列
# )
train_args = TrainingArguments(
    output_dir=finetuned_model_path,
    num_train_epochs=3,  # 进一步减少轮数
    per_device_train_batch_size=1,  # 关键：将批次从4降为1（显存优先）
    gradient_accumulation_steps=16,  # 用累积步数弥补小批次，保持总更新量
    save_steps=50,
    save_total_limit=2,
    report_to="swanlab",
    logging_dir="./logs",
    logging_steps=5,
    save_strategy="steps",
    gradient_checkpointing=True,  # 重新开启梯度检查点（节省30%+显存）
    remove_unused_columns=False,
    max_steps=200,  # 减少总步数
    learning_rate=5e-4,
    fp16=True,  # 保持FP16混合精度
    # 新增：梯度裁剪（防止梯度爆炸同时减少显存波动）
    max_grad_norm=1.0,
)

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False,#是否使用MLM(标准语言模型任务)
)

trainer = Trainer(
    model=model,
    args=train_args,
    data_collator=data_collator,
    train_dataset=dataset,
)

#开始训练
trainer.train()